package com.leiyk.rateLimiter.aop;

import com.leiyk.rateLimiter.annotation.SimpleRateLimiter;
import com.leiyk.rateLimiter.constants.RedisConstant;
import com.leiyk.rateLimiter.exception.RateLimiterException;
import jakarta.annotation.Resource;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.util.List;


/**
 * @Description: 注解aop实现类
 * @Author: leiyk6
 * @Date: 2024/7/3 10:08
 */
@Aspect
@Component
public class SimpleRateLimiterAop {
    @Pointcut("@annotation(com.leiyk.rateLimiter.annotation.SimpleRateLimiter)")
    void pointCut(){}

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    private static final String scr =
            "local key = KEYS[1]\n" +
            "local time = ARGV[1]\n" +
            "local max_count = tonumber(ARGV[2])\n" +
            "local curCount = redis.call('INCR', key)\n" +
            "if curCount == 1 then\n   " +
            "redis.call('EXPIRE',key, time)\n" +
            "end\n" +
            "if curCount > max_count then\n   " +
            "return 0\n" +
            "end\n" +
            "return 1";


    @Around("pointCut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        SimpleRateLimiter annotation = signature.getMethod().getAnnotation(SimpleRateLimiter.class);
        int maxCount = annotation.maxCount();
        int time = annotation.time();
        String msg = annotation.msg();

        String key = RedisConstant.SIMPLE_RATE_LIMITER_PREFIX + signature;

        System.out.println(scr);
        Long val = stringRedisTemplate.execute(new DefaultRedisScript<>(scr, Long.class), List.of(key), String.valueOf(time), String.valueOf(maxCount));
        if(val == 0L){
            System.out.println("拒绝服务");
            throw RateLimiterException.rejectRequest();
        }
        return joinPoint.proceed();

    }


}
